import cv2
import time
import numpy
import torch
import config as cfg
import pickle
from multiprocessing import Pool
from models.facenet.facenet import Facenet
from copy import copy
from utils import plot_one_box
from numpy import load, argmin, array, zeros
from face_recognition import face_distance, face_locations, face_landmarks
from threading import Thread


Image = None

def send(name):
    print(name)

def detect_img(img, codeing_model, knn, thres=1.5, size_thres=400):
    names = []
    # 减小输入尺寸，加快识别速度， 但牺牲了对小人脸的检测能力
    small_img = cv2.resize(img, (0, 0), fx=0.25, fy=0.25)[:, :, ::-1]
    boxes = face_locations(small_img, model='hog')
    boxes = array(boxes)*4  # 将检测框映射回原图

    # 人脸裁剪
    max_box = None
    for box in boxes:  # box => top, right, bottom, left
        x, y = box[2]-box[0], box[1]-box[3]
        value = x*y
        if value > size_thres:
            max_box = box
            size_thres = value
    if max_box is None:
        name = None
    else:
        face = numpy.expand_dims(img[max_box[0]:max_box[2], max_box[3]:max_box[1]], 0)

        # 人脸编码
        feature = None
        with torch.no_grad():
            face = torch.from_numpy(face.transpose(0, 3, 1, 2)).type(torch.FloatTensor)
            feature = codeing_model(face).detach().numpy()

        # 判断是否为陌生人
        # closest_distances = knn.kneighbors(feature, n_neighbors=1)
        # print(closest_distances[0][0][0])
        # print(thres)
        # 映射到name表， 同时绘制矩形框
        name = knn.predict(feature)[0]#  if closest_distances[0][0][0] <= thres else 'unknow'
        img = plot_one_box(array(max_box), img, label=name, line_thickness=1)
    return img, name


def detect_video(codeing_model, video_path=0):
    # 配置摄像头
    def get_stream():
        capture = cv2.VideoCapture(video_path)
        global IMAGE
        ref, frame = capture.read()
        while True:
            _, IMAGE = capture.read()
    thread = Thread(target=get_stream).start()
    person = [None]*10
    capture = cv2.VideoCapture(video_path)
    knn = pickle.load(open('data\knn.pickle', 'rb'))
    while True:
        t1 = time.time()
        ref, frame = capture.read()
        if not ref:
            break
        
        # 从图片识别人脸，返回框选后的图片和对应的人名列表
        img, name = detect_img(frame, codeing_model, knn)
        if name:
            dic = {}
            person = person[1:] + [name]
            for p in person:
                if p:
                    dic[p] = dic[p]+1 if p in dic.keys() else 1
            keys = []
            for k in dic.keys():
                keys.append(k)
            result = keys[numpy.argmax(dic.values())]

            send(result)
        else:
            person = [None]*10

        cv2.imshow('video', img)
        cv2.waitKey(1)



if __name__=='__main__':
    # 加载模型
    path = 'models/facenet/facenet_mobilenet.pth'
    model = Facenet(mode='predict').eval()
    state_dict = torch.load(path, map_location='cpu')
    model.load_state_dict(state_dict, strict=False)

    # 预测视频流
    # detect_video(codeing_model=model, video_path='rtmp://rtmp01open.ys7.com:1935/v3/openlive/E95461747_1_1?expire=1669346598&id=385752649638678528&t=52f714afce0613af4759f09e2dde6ea294cbc998e1f218e2d42aa0d8462649de&ev=100')
    detect_video(codeing_model=model, video_path=0)